package com.ruoyi.chat.store;

import com.ruoyi.chat.domain.ChatMessages;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.LinkedList;
import java.util.List;
import java.util.TimeZone;

/**
 * @author WaKy
 * @version 1.0
 */
@Component
public class MongoChatMemoryStore implements ChatMemoryStore {

    @Autowired
    private MongoTemplate mongoTemplate;

    /**
     * 根据会话ID获取聊天消息列表
     *
     * @param memoryId 会话ID
     * @return 该会话的消息列表
     */
    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
        if (chatMessages == null) {
            return new LinkedList<>();
        }
        return ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
    }

    /**
     * 更新指定会话的消息列表
     *
     * @param memoryId 会话ID
     * @param messages 新的消息列表
     */
    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        // 将messages中的SystemMessage过滤掉(没用并且浪费存储空间)
        messages.removeIf(message -> message instanceof SystemMessage);
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        Update update = new Update();
        update.set("content", ChatMessageSerializer.messagesToJson(messages));
        // 添加时间
        // 如果当前memoryId对应的文档不存在，则更新时间，否则不更新时间
        if (getMessages(memoryId).isEmpty()) {
            SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
            // 指定时区为上海
            sdf.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));
            update.set("chatTime", sdf.format(new Date()));
        }

        // 根据query条件能查询出文档，则修改文档，查询不出就新增文档
        mongoTemplate.upsert(query, update, ChatMessages.class);
    }

    /**
     * 删除指定会话的所有消息
     *
     * @param memoryId 会话ID
     */
    @Override
    public void deleteMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        mongoTemplate.remove(query, ChatMessages.class);
    }
}
